#include <tshark_database.h>

namespace Tshark
{
    bool TsharkDataBase::openDataBase(const std::string &dbName)
    {
        sqlite3 *dbPtr;
        int rc = sqlite3_open(dbName.c_str(), &dbPtr);
        if (dbPtr == nullptr || rc != SQLITE_OK)
        {
            LOG_F(ERROR, "%s", sqlite3_errmsg(dbPtr));
            return false;
        }
        db = std::unique_ptr<sqlite3, std::function<int(sqlite3 *)>>(dbPtr,
                                                                     [](sqlite3 *ptr)
                                                                     { return sqlite3_close(ptr); });
        this->createPacketTable();
        return true;
    }

    bool TsharkDataBase::storePackets(const std::vector<std::shared_ptr<DataType::Packet>> &packets)
    {
        if (packets.empty())
        {
            return true;
        }
        // 事务有满足 ACID 的4个原则
        char *errMsg;
        int rc = sqlite3_exec(db.get(), "BEGIN TRANSACTION;", nullptr, nullptr, &errMsg);
        if (rc != SQLITE_OK)
        {
            LOG_F(ERROR, "%s", errMsg);
            sqlite3_free(errMsg);
            return false;
        }

        std::string insertSQL = R"(INSERT INTO t_packets (
                frame_number, time, src_ip, src_location, src_port, dst_ip, dst_location, dst_port,
                protocol, info, cap_len, len, src_mac, dst_mac, file_offset
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
        )";
        sqlite3_stmt *stmt;
        if (sqlite3_prepare_v2(db.get(), insertSQL.c_str(), -1, &stmt, nullptr) != SQLITE_OK)
        {
            LOG_F(ERROR, "%s", sqlite3_errmsg(db.get()));
            throw std::runtime_error("Failed to prepare insert statement");
        }

        bool hasError = false;
        for (const auto &packet : packets)
        {
            sqlite3_bind_int(stmt, 1, packet->frame_number);
            sqlite3_bind_double(stmt, 2, packet->time);
            sqlite3_bind_text(stmt, 3, packet->src_ip.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_text(stmt, 4, packet->src_location.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_int(stmt, 5, packet->src_port);
            sqlite3_bind_text(stmt, 6, packet->dst_ip.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_text(stmt, 7, packet->dst_location.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_int(stmt, 8, packet->dst_port);
            sqlite3_bind_text(stmt, 9, packet->protocol.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_text(stmt, 10, packet->info.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_int(stmt, 11, packet->cap_len);
            sqlite3_bind_int(stmt, 12, packet->len);
            sqlite3_bind_text(stmt, 13, packet->src_mac.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_text(stmt, 14, packet->dst_mac.c_str(), -1, SQLITE_STATIC);
            sqlite3_bind_int(stmt, 15, packet->file_offset);
            if (sqlite3_step(stmt) != SQLITE_DONE)
            {
                hasError = true;
                LOG_F(ERROR, "Failed to execute insert statement: %s", sqlite3_errmsg(db.get()));
                break;
            }
            sqlite3_reset(stmt); // 重置,进行下一次绑定
        }
        if (!hasError)
        {
            // 提交事务
            if (sqlite3_exec(db.get(), "COMMIT;", nullptr, nullptr, nullptr) != SQLITE_OK)
            {
                hasError = true;
            }
        }
        else
        {
            // 出错事务回滚
            sqlite3_exec(db.get(), "ROLLBACK TRANSACTION;", nullptr, nullptr, &errMsg);
        }
        // 释放语句
        sqlite3_finalize(stmt);
        return !hasError;
    }

    bool TsharkDataBase::queryPackets(std::vector<std::shared_ptr<DataType::Packet>> &packetList)
    {
        sqlite3_stmt *stmt = nullptr;
        std::string querySQL = "select * from t_packets";
        if (sqlite3_prepare_v2(db.get(), querySQL.c_str(), -1, &stmt, nullptr) != SQLITE_OK)
        {
            LOG_F(ERROR, "Failed to prepare statement: ");
            return false;
        }
        while (sqlite3_step(stmt) != SQLITE_DONE)
        {
            std::shared_ptr<DataType::Packet> packet = std::make_shared<DataType::Packet>();
            packet->frame_number = sqlite3_column_int(stmt, 0);
            packet->time = sqlite3_column_double(stmt, 1);
            packet->src_ip = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 2));
            packet->src_location = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 3));
            packet->src_port = sqlite3_column_int(stmt, 4);
            packet->dst_ip = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 5));
            packet->dst_location = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 6));
            packet->dst_port = sqlite3_column_int(stmt, 7);
            packet->protocol = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 8));
            packet->info = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 9));
            packet->cap_len = sqlite3_column_int(stmt, 10);
            packet->len = sqlite3_column_int(stmt, 11);
            packet->src_mac = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 12));
            packet->dst_mac = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 13));
            packet->file_offset = sqlite3_column_int(stmt, 14);
            packetList.push_back(packet);
        }
        sqlite3_finalize(stmt);
        return true;
    }

    bool TsharkDataBase::queryPackets(const std::string &ip, uint16_t port, std::vector<std::shared_ptr<DataType::Packet>> &packets)
    {
        sqlite3_stmt *stmt = nullptr;
        std::string querySQL = R"(select * from t_packets where (src_ip = ? or dst_ip = ?) and (src_port = ? or dst_port = ?) and src_location GLOB "*";)";
        if (sqlite3_prepare_v2(db.get(), querySQL.c_str(), -1, &stmt, nullptr) != SQLITE_OK)
        {
            LOG_F(ERROR, "Failed to prepare statement: %s", sqlite3_errmsg(db.get()));
            return false;
        }
        sqlite3_bind_text(stmt, 1, ip.c_str(), -1, nullptr);
        sqlite3_bind_text(stmt, 2, ip.c_str(), -1, nullptr);
        sqlite3_bind_int(stmt, 3, port);
        sqlite3_bind_int(stmt, 4, port);
        while (sqlite3_step(stmt) != SQLITE_DONE)
        {
            std::shared_ptr<DataType::Packet> packet = std::make_shared<DataType::Packet>();
            packet->frame_number = sqlite3_column_int(stmt, 0);
            packet->time = sqlite3_column_double(stmt, 1);
            packet->src_ip = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 2));
            packet->src_location = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 3));
            packet->src_port = sqlite3_column_int(stmt, 4);
            packet->dst_ip = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 5));
            packet->dst_location = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 6));
            packet->dst_port = sqlite3_column_int(stmt, 7);
            packet->protocol = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 8));
            packet->info = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 9));
            packet->cap_len = sqlite3_column_int(stmt, 10);
            packet->len = sqlite3_column_int(stmt, 11);
            packet->src_mac = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 12));
            packet->dst_mac = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 13));
            packet->file_offset = sqlite3_column_int(stmt, 14);
            packets.push_back(packet);
        }
        return true;
    }

    bool TsharkDataBase::queryPackets(const QueryConditon &conditon, std::vector<std::shared_ptr<DataType::Packet>> &packets)
    {
        std::string querySQL = SQL::PacketSQL::buildPacketQuerySQL(conditon);

        sqlite3_stmt *stmt;
        if (sqlite3_prepare_v2(db.get(), querySQL.c_str(), -1, &stmt, nullptr)!=SQLITE_OK)
        {
            LOG_F(ERROR, "sql:%s sqlite3_prepare_v2 error:%s", querySQL.c_str(), sqlite3_errmsg(db.get()));
            return false;
        }
        while (sqlite3_step(stmt) == SQLITE_ROW)
        {
            std::shared_ptr<DataType::Packet> packet = std::make_shared<DataType::Packet>();
            packet->frame_number = sqlite3_column_int(stmt, 0);
            packet->time = sqlite3_column_double(stmt, 1);
            packet->src_ip = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 2));
            packet->src_location = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 3));
            packet->src_port = sqlite3_column_int(stmt, 4);
            packet->dst_ip = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 5));
            packet->dst_location = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 6));
            packet->dst_port = sqlite3_column_int(stmt, 7);
            packet->protocol = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 8));
            packet->info = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 9));
            packet->cap_len = sqlite3_column_int(stmt, 10);
            packet->len = sqlite3_column_int(stmt, 11);
            packet->src_mac = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 12));
            packet->dst_mac = reinterpret_cast<const char *>(sqlite3_column_text(stmt, 13));
            packet->file_offset = sqlite3_column_int(stmt, 14);
            packets.push_back(packet);
        }
        return true;
    }

    TsharkDataBase::~TsharkDataBase()
    {
        // 交给 std::unique_ptr 释放资源
        // if (db != nullptr)
        // {
        //     sqlite3_close(db.get());
        // }
    }

    bool TsharkDataBase::createPacketTable()
    {
        // 检查表是否存在，若不存在则创建
        std::string createTableSQL = R"(
            CREATE TABLE IF NOT EXISTS t_packets (
                frame_number INTEGER PRIMARY KEY,
                time REAL,
                src_ip TEXT,
                src_location TEXT,
                src_port INTEGER,
                dst_ip TEXT,
                dst_location TEXT,
                dst_port INTEGER,
                protocol TEXT,
                info TEXT,
                cap_len INTEGER,
                len INTEGER,
                src_mac TEXT,
                dst_mac TEXT,
                file_offset INTEGER
            );
        )";
        char *errMsg;
        int rc = sqlite3_exec(db.get(), createTableSQL.c_str(), nullptr, nullptr, &errMsg);
        if (rc != SQLITE_OK)
        {
            LOG_F(ERROR, "%s", errMsg);
            sqlite3_free(errMsg);
            return false;
        }
        LOG_F(INFO, "create table t_packets successfully!");
        return true;
    }
};